import numpy as np
import matplotlib.pyplot as plt
import time
from matplotlib import cm
from scipy.interpolate import griddata
import matplotlib.colors as colors
from mpl_toolkits.axes_grid1 import make_axes_locatable
from matplotlib.patches import Rectangle
import io
import os
import json
from scipy import signal
import pandas as pd

temp = []
IPValue = {}
fig, ax = plt.subplots(figsize=(15, 7.5))


def ERTuiArray(ABMN, a1_dtype, a1_shape, Aspacing):
    abmn = ABMN.encode('ISO-8859-1')
    ABMN1 = np.fromstring(abmn, dtype=getattr(np, a1_dtype)).reshape(eval(a1_shape))
    # print(print(ABMN1[0,:]))
    return ABMN1


def AnalyzeJonsJson(json_str, ns,nss):
    # 将JSON字符串转换为Python对象
    data = json.loads(json_str)
    # ReadData(data)
    io = ERTuiSection(data, int(ns), ns,nss)
    return io


def ERTuiSectionNs(json_str):
    data = json.loads(json_str)
    print("json abmn =======================================")
    print(data)
    # NS = np.sum(range(1,int(data['Max_n_spacing'])+1)) + int(data['Max_n_spacing'])*((int(data['num_electrodes'])-3)-int(data['Max_n_spacing']))      #总排列数
    NS = len(ERTuiArray(data['ABMN'], data['dtype'], data['shape'], data['Aspacing']))
    print(NS)
    return int(NS)


# 排列图示
# def ERTuiSection(data,i):
#     AI = int(i)
#     print("-----------------------AI")
#     print(AI)
#     ABMN = ERTuiArray(data['ABMN'], data['dtype'], data['shape'], data['Aspacing'])
#
#     #plt.ion()  # 打开交互模式
#
#     Data = ReadData(data, ABMN[AI, :])
#
#     ###4. 处理数据
#
#     IPValue[AI] = ComputeRhoIP(Data, data, AI)
#     print("---------------------IPValue")
#     print(IPValue)
#     ###5. 画出图形
#     io = ERTPlot(IPValue, ABMN, int(data['Aspacing']))
#     return io
def ERTuiSection(data, AI, ns,nss):
    # NS = np.sum(range(1,int(data['Max_n_spacing'])+1)) + int(data['Max_n_spacing'])*((int(data['num_electrodes'])-3)-int(data['Max_n_spacing']))      #总排列数
    ###3. 读取CVS数据及系统参数
    '''
        手部首先需要向主机发送数据采集的命令，规定采集的参数，然后主机才能返回数据。这里假设命令发送成果，数据成功上传。
        数据按每次叠加上传，例如我们设置命令为叠加5次，每次波形长度为8，那么手部会收到8个完整波形后再做相应的处理。
        这8个波形的数据在数据文件中是一个数据块，有相应的头文件与之对应。
    '''
    ABMN = ERTuiArray(data['ABMN'], data['dtype'], data['shape'], data['Aspacing'])
    # IPValue = {}
    plt.ion()  # 打开交互模式
    # for AI in range(NS):
    Data = ReadData(data, ABMN[AI, :])
    print("----------------------ABMN[AI,:]------------")
    print(ABMN[AI, :])

    ###4. 处理数据

    IPValue[AI] = ComputeRhoIP(Data, data, AI)
    # print(IPValue)
    ###5. 画出图形
    io = ERTPlot(IPValue, ABMN, int(data['Aspacing']), ns,nss)
    return io


def ReadData(data, ABMN):
    '''
        程序首先发送参数到主机，然后接收主机返回的数据和其它参数
        :param SysInfo: 系统信息
        :param SurveyInfo: 项目信息
        :param ABMN: 发射和接收电极位置
        :return: SysInfo, SurveyInfo, Data
    '''
    ###发送给主机的参数
    TXFrq = int(data['TXFrq'])  # 发射频率
    TXDuty = int(data['TXDuty'])  # 发射占空比
    CycleNO = int(data['CycleNO'])  # 一次叠加波形数
    StackNO = int(data['StackNO'])  # 叠加次数
    TXLoc = ABMN[0:2]  # 发射电极位置
    TXMode = 'V'  # 发射模式，V=最大电压发射，I=最大电流发射
    RXLoc = ABMN[2:]  # 接收电极位置
    Filter = True  # 50Hz滤波是否打开
    Gain = True  # 自动增益是否打开

    ####发送命令给主机
    DataSTK = {}
    # SurveyInfo['STKCTDN'] = 0                               # 叠加倒计数 初始化
    for I in range(StackNO):
        Data = SendCMD(TXFrq, TXDuty, CycleNO, TXLoc, TXMode, RXLoc, Filter, Gain, data['path'])
        if len(Data):  # 成功采集到数据
            # 将数据波形叠加平均
            SR = int(data['RXSR'])
            temp = np.zeros((int(SR / TXFrq), 2))
            for J in range(CycleNO):
                stp = int(J * SR / TXFrq)
                endp = int((J + 1) * SR / TXFrq)
                temp = temp + Data[stp:endp, :]
            DataSTK[I] = temp / CycleNO
        else:
            print('数据采集出错，请检测')
            return
    return DataSTK


def SendCMD(TXFrq, TXDuty, CycleNO, TXLoc, TXMode, RXLoc, Filter, Gain, path):
    '''
        发送命令到主机，返回系统信息，数据
        :param TXFrq:
        :param TXDuty:
        :param CycleNO:
        :param StackNO:
        :param TXLoc:
        :param TXMode:
        :param RXLoc:
        :param Filter
        :param Gain
        :return: SysInfo, Data
    '''
    SysInfo = {}
    SysInfo['SysBatV'] = 12.5  # 仪器电池电压
    SysInfo['SysTemp'] = 30  # 仪器温度
    SysInfo['RXSR'] = 1000  # 采样频率
    SysInfo['Gain'] = 0  # 增益关
    SysInfo['Filter'] = 1  # 50Hz陷波开
    SysInfo['CRES'] = 0.8  # 接地电阻

    ###这里采用读取一个cvs文件的方式模拟仪器返回的一个数据
    CurrentDir = os.path.dirname(os.path.abspath(__file__))  # current direcotry
    os.chdir(CurrentDir)
    # filepath = path
    # print("_______________________filepath_______" + str(filepath))
    # Data = Read_cvs(filepath)
    # print("-=-=-=-=- Read_cvs_data = "+str(Data))
    # 实际数据采集中，我们只需要保存发射电流和接收电压的值。现在仪器采集的数据和我们发送给仪器的命令是不同的，因此这里采用合成的方式形成采集的数据
    # 上述的数据读取部分可供参考
    # !!!!!以下为数据模拟过程，不需要编写到手部里
    NFFT = SysInfo['RXSR'] * CycleNO / TXFrq  # number of points
    SR = SysInfo['RXSR']  # sampling rate
    A = 1  # mag
    Phs = 0  # phase
    SF = TXFrq  # square wave frequency
    Delta_T = 1 / SR  # sampling interval
    t = np.arange(0, NFFT, 1) * Delta_T  # Time vector
    TX_I = A * signal.square(2 * np.pi * SF * t + Phs) + \
           A * signal.square(2 * np.pi * SF * t + Phs + np.pi / 2)  # TX_I wave
    RX_V = A * signal.square(2 * np.pi * SF * t + Phs) + \
           A * signal.square(2 * np.pi * SF * t + Phs + np.pi / 2)  # RX_V wave
    Decay = 2 * np.exp(-16 * np.pi * SF * np.arange(0, SR / TXFrq / 4) * Delta_T)
    temp = np.zeros((int(SR / TXFrq),))
    temp[int(SR / TXFrq / 4):int(SR / TXFrq / 2)] = Decay
    temp[int(3 * SR / TXFrq / 4):int(SR / TXFrq)] = -Decay
    # add Decay to RX_V to simulate a response
    for I in range(CycleNO):
        stp = int(I * SR / TXFrq)
        endp = int((I + 1) * SR / TXFrq)
        RX_V[stp:endp] = RX_V[stp:endp] + temp
    Data = np.vstack((TX_I, RX_V)).T
    return Data


def Read_cvs(fileloc):
    '''This routine reads the data inside cvs files.
       The CVS file contains column name and data.
    '''
    try:
        # read in data as data frame
        data = pd.read_csv(fileloc)
        # change TX_V to V
        data['TX_V'] = data['TX_V'] * 100
        return data
    except Exception as e:
        print('Read txt', e)


def ComputeRhoIP(Data, data, AI):
    '''
        计算电阻，极化率，及画图的位置.
        此算法暂时只针对偶极-偶极排列
        :param Data:
        :param SysInfo:
        :param SurveyInfo:
        :param ABMN:
        :return: IPValues
    '''
    ABMN1 = ERTuiArray(data['ABMN'], data['dtype'], data['shape'], data['Aspacing'])
    ABMN = ABMN1[AI, :]
    SR = int(data['RXSR'])  # 数据点
    TXFrq = int(data['TXFrq'])  # 发射频率
    StackNO = int(data['StackNO'])  # 叠加次数  这三个都是app填写的
    Aspacing = int(data['Aspacing'])
    TXA = (ABMN[0] - 1) * Aspacing  # Aspacing 点距 也是app 填写的
    TXB = (ABMN[1] - 1) * Aspacing
    RXM = (ABMN[2] - 1) * Aspacing
    RXN = (ABMN[3] - 1) * Aspacing
    IPValues = {}
    IPValues['TXA'] = TXA
    IPValues['TXB'] = TXB
    IPValues['RXM'] = RXM
    IPValues['RXN'] = RXN
    IPValues['X'] = (RXM - TXB) / 2 + TXB  # 画图横坐标位置
    IPValues['Z'] = (RXM - TXB) / 2 + Aspacing / 2  # 画图垂向坐标位置
    IPValues['n'] = ABMN[2] - ABMN[1]  # 隔离系数

    # 叠加处理数据
    temp = Data[0]  # 将第一个叠加数据暂存起来
    for BI in range(1, StackNO):
        temp = temp + Data[BI]  # 叠加数据
    temp = temp / StackNO  # 平均数据
    NP = int(SR / TXFrq)  # 一个周期的数据点数
    NPhlf = int(NP / 2)  # 半个周期数据点数
    tmp = (temp[0:NPhlf, :] - temp[NPhlf:, :]) / 2  # 将后半个周期和前半个周期平均
    # print("---------------- tmp ="+str(tmp))
    TX_I = tmp[:, 0]  # 发射电流数据
    # print("---------------- TX_I ="+str(TX_I))
    RX_V = tmp[:, 1]  # 接收电压数据
    VpPnts = int(0.0625 * NP)  # 数据点数用于计算电流和电压
    TXI = np.sum(TX_I[int(NPhlf / 2 - VpPnts):int(NPhlf / 2)]) / VpPnts  # 发射电流值
    Vp = np.sum(RX_V[int(NPhlf / 2 - VpPnts):int(NPhlf / 2)]) / VpPnts  # 接收电压值
    IPValues['TXI'] = TXI
    IPValues['RXV'] = Vp
    Mpts = int(0.08125 * NP)
    Mst = int(0.05625 * NP) + 4
    IP = 1 / TXFrq / NP * 1.87 / Vp * np.sum(RX_V[int(NPhlf + Mst):int(NPhlf + Mst + Mpts)]) / 10
    IPValues['IP'] = IP
    Rho = Vp / TXI * np.pi * IPValues['n'] * (IPValues['n'] + 1) * (IPValues['n'] + 2) * Aspacing
    IPValues['Rho'] = Rho
    return IPValues


def ERTPlot(IPValues, ABMN, Aspacing, NS,NSS):
    print("___________________ERTPlot__________________11111")
    '''
        本程序画出高密度电阻率法的等值线图和排列图
        :param IPValues:
        :return: NONE
    '''
    ####这里画出两个图，一个是排列图，利用红色圆圈表示发射，用黑色圆表示接收
    # 最后一个IP的值
    # 获取相应的画图参数
    # print("-------------------------IPValues---------------------------")
    # print(IPValues)
    NIP = len(IPValues)

    # #####计算图的尺寸
    Min_X = np.min(ABMN[:, 0] - 1) * Aspacing
    Max_X = np.max(ABMN[:, 3] - 1) * Aspacing
    Min_Z = np.min(ABMN[:, 2] - ABMN[:, 1]) * Aspacing / 2 + Aspacing / 2
    Max_Z = np.max(ABMN[:, 2] - ABMN[:, 1]) * Aspacing / 2 + Aspacing / 2

    colormap = cm.get_cmap('jet')
    plt.xlim([Min_X, Max_X])
    plt.ylim([0, Max_Z + 10])
    #plt.annotate(format(rho, '.0f'), xy=(X, Z), fontsize=8, ha='center', va='center')
    plt.gca().invert_yaxis()
    plt.gca().set_xlabel('Array Location(m)', size=14)
    plt.gca().set_ylabel('n-spacing', size=14)
    plt.gca().xaxis.set_label_position('top')
    plt.gca().xaxis.tick_top()
    plt.gca().set_aspect('equal')
    plt.yticks(np.insert(np.arange(Min_Z, Max_Z + Aspacing / 2, Aspacing / 2), 0, 0))
    plt.xticks(np.arange(Min_X, Max_X, Aspacing))

    print("-----------------------------------NLAB--------------------------------------------")
    print(NIP)
    print(NSS)
    if NIP == NSS:
        plt.clf()
        ax.axes.get_xaxis().set_visible(False)
        CX = np.zeros((NSS,))
        CZ = np.zeros((NSS,))
        Rho = np.zeros((NSS,))
        for I in range(NSS):
            currentIP = IPValues[I]
            CX[I] = currentIP['X']
            CZ[I] = currentIP['Z']
            Rho[I] = currentIP['Rho']
        #网格化
        xv, zv = np.meshgrid(np.linspace(Min_X, Max_X, 100), np.linspace(Min_Z, Max_Z, 100))
        #插值
        rho_av = griddata((CX, CZ), Rho, (xv, zv))
        cs = plt.contourf(xv, zv, rho_av, levels=np.logspace(1,5,40), cmap=cm.jet,norm = colors.LogNorm())
        cs.set_clim(vmin=10, vmax=10000)
        # plt.gca().clabel(cs, fmt='%2.1f', colors='k', fontsize=8, inline=True,inline_spacing=0)
        plt.gca().invert_yaxis()
        plt.gca().set_xlabel('Array Location(m)', size=14)
        plt.gca().set_ylabel('n-spacing', size=14)
        plt.gca().xaxis.set_label_position('top')
        plt.gca().xaxis.tick_top()
        plt.gca().set_aspect('equal')
        plt.yticks(np.insert(np.arange(Min_Z, Max_Z + Aspacing / 2, Aspacing / 2), 0, 0))
        plt.xticks(np.arange(Min_X, Max_X, Aspacing))
        labels = [item.get_text() for item in ax.get_yticklabels()]
        NLAB = len(labels)
        labs = np.insert(np.arange(int(np.min(ABMN[:, 2] - ABMN[:, 1])), \
                                   int(np.max(ABMN[:, 2] - ABMN[:, 1]) + 1), 1), 0, 0)
        if NLAB == len(labs):
            ax.set_yticklabels(labs)
        divider = make_axes_locatable(ax)
        cax = divider.append_axes('right', size="2%", pad=0.1)
        sm = plt.cm.ScalarMappable(norm=colors.LogNorm(), cmap=colormap)
        sm.set_clim(vmin=10, vmax=10000)
        cbar = plt.colorbar(sm, cax=cax)
        cbar.set_label('Apparent Resistivity(Ohm-m)', rotation=270, labelpad=20)
        ax.spines['bottom'].set_visible(False)
        plt.show()
        #cbar.remove()
        f = io.BytesIO()
        plt.savefig(f, format="png")
        print("___________________ERTPlot__________________33333")
        return f.getvalue()
    f = io.BytesIO()
    plt.savefig(f, format="png")
    print("___________________ERTPlot__________________33333")
    return f.getvalue()